package org.shx.springboot.ollama.ai;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ArrayNode;
import com.fasterxml.jackson.databind.node.ObjectNode;
import org.springframework.http.MediaType;
import org.springframework.web.reactive.function.BodyInserters;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Flux;

public class OllamaToolTest {

    public static void main(String[] args) {
        new OllamaToolTest().testToolCall().blockLast();
    }

    public Flux<String> testToolCall() {
        // 1. 创建ObjectMapper用于构建JSON
        ObjectMapper mapper = new ObjectMapper();

        // 2. 构建工具定义
        ObjectNode tool1 = mapper.createObjectNode();
        tool1.put("name", "getTableNameAndComment");
        tool1.put("description", "查询所有表名和注释");
        ObjectNode params1 = mapper.createObjectNode();
        params1.put("type", "object");
        params1.set("properties", mapper.createObjectNode());
        params1.set("required", mapper.createArrayNode());
        tool1.set("parameters", params1);

        ArrayNode toolsArray = mapper.createArrayNode();
        toolsArray.add(tool1);

        // 3. 构建消息
        ArrayNode messagesArray = mapper.createArrayNode();

        ObjectNode systemMsg = mapper.createObjectNode();
        systemMsg.put("role", "system");
        systemMsg.put("content", "必须使用工具回答，直接返回工具调用的JSON");
        messagesArray.add(systemMsg);

        ObjectNode userMsg = mapper.createObjectNode();
        userMsg.put("role", "user");
        userMsg.put("content", "查询所有表名和注释");
        messagesArray.add(userMsg);

        // 4. 构建完整请求体
        ObjectNode requestBody = mapper.createObjectNode();
        requestBody.put("model", "qwen3:8b");
        requestBody.set("messages", messagesArray);
        requestBody.set("tools", toolsArray);
        requestBody.put("tool_choice", "auto");
        requestBody.put("stream", true);

        // 5. 创建WebClient并添加拦截器
        WebClient webClient = WebClient.builder()
                .baseUrl("http://172.16.20.61:11434")
                .filter((request, next) -> {
                    // 打印请求信息
                    System.out.println("请求方法: " + request.method());
                    System.out.println("请求URL: " + request.url());

                    // 打印请求体（Spring AI 1.0.0兼容方式）
                    if (request.body() != null) {
                        try {
                            // 将请求体转换为字符串
                            String bodyStr = new ObjectMapper()
                                    .writeValueAsString(request.body());
                            System.out.println("请求体: " + bodyStr);
                        } catch (Exception e) {
                            System.err.println("无法解析请求体: " + e.getMessage());
                        }
                    }
                    return next.exchange(request);
                })
                .build();

        // 6. 发送请求并处理响应
        return webClient.post()
                .uri("/api/chat")
                .contentType(MediaType.APPLICATION_JSON)
                .body(BodyInserters.fromValue(requestBody))
                .exchangeToFlux(response -> {
                    // 处理错误状态码
                    if (response.statusCode().is4xxClientError()) {
                        return response.bodyToFlux(String.class)
                                .flatMap(error -> {
                                    System.err.println("Ollama错误响应: " + error);
                                    return Flux.error(new RuntimeException("Ollama返回4xx错误: " + error));
                                });
                    }
                    // 处理正常响应
                    return response.bodyToFlux(String.class);
                });
    }

    /*{"error":"registry.ollama.ai/library/llama3:latest does not support tools"} 可知，当前使用的 llama3:latest 模型不支持工具调用功能。*/
}

